from src.dataset import FudanDrowsyDataset
import torch
import torch.utils.data
from torch import nn
import numpy as np
from torch.nn import functional as F
from src import const
from src.utils import parse_args_and_merge_const, get_train_test_df, Evaluator, get_train_test_df_with_in_use
from tensorboardX import SummaryWriter
import os


if __name__ == '__main__':
    parse_args_and_merge_const()
    if os.path.exists('models') is False:
        os.makedirs('models')

    if const.USE_SAMPLE_FILTERED:
        train_df, test_df = get_train_test_df_with_in_use()
        train_df_no_filtered = train_df
        train_df = train_df[train_df['in_use'] == True]
    else:
        train_df, test_df = get_train_test_df()
        train_df_no_filtered = train_df
    train_dataset = FudanDrowsyDataset(train_df, mode=const.TRAIN_DATASET_MODE)
    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=const.BATCH_SIZE, shuffle=True, num_workers=4)
    train_dataset_unshuffle = FudanDrowsyDataset(train_df_no_filtered, mode=const.VAL_DATASET_MODE)
    train_dataloader_unshuffle = torch.utils.data.DataLoader(train_dataset_unshuffle, batch_size=const.VAL_BATCH_SIZE, shuffle=False, num_workers=4)
    test_dataset = FudanDrowsyDataset(test_df, mode=const.VAL_DATASET_MODE)
    test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=const.VAL_BATCH_SIZE, shuffle=False, num_workers=4)

    net = const.USE_NET()
    net = net.to(const.device)

    learning_rate = const.LEARNING_RATE
    optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)

    writer = SummaryWriter(const.TRAIN_DIR)

    total_step = len(train_dataloader)
    step = 0
    criterion = nn.CrossEntropyLoss()
    for epoch in range(const.NUM_EPOCH):
        net.train()
        for i, sample in enumerate(train_dataloader):
            step += 1
            for key in sample:
                sample[key] = sample[key].to(const.device)
            output = net(sample['image'])
            loss = criterion(output, sample['label'])

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if (i + 1) % 10 == 0:
                writer.add_scalar('loss', loss.item(), step)
                writer.add_scalar('learning_rate', learning_rate, step)
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                      .format(epoch + 1, const.NUM_EPOCH, i + 1, total_step, loss.item()))

        print('Saving Model....')
        net.set_buffer('step', step)
        torch.save(net.state_dict(), 'models/' + const.MODEL_NAME)
        print('OK. Now evaluate..')

        net.eval()  # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)
        with torch.no_grad():
            evaluator = Evaluator(test_df)
            for i, sample in enumerate(test_dataloader):
                for key in sample:
                    sample[key] = sample[key].to(const.device)
                output = net(sample['image'])
                evaluator.add(sample, output)
            ret = evaluator.evaluate()
            for key in ret:
                print(key, ret[key])
                writer.add_scalar('metrics/' + key, ret[key], step)
        with torch.no_grad():
            evaluator = Evaluator(train_df_no_filtered)
            for i, sample in enumerate(train_dataloader_unshuffle):
                for key in sample:
                    sample[key] = sample[key].to(const.device)
                output = net(sample['image'])
                evaluator.add(sample, output)
            ret = evaluator.evaluate()
            for key in ret:
                print('train_' + key, ret[key])
                writer.add_scalar('monitor/train_' + key, ret[key], step)
        # learning rate decay
        learning_rate *= const.LEARNING_RATE_DECAY
        optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
